
import os
from posixpath import pardir
import time
import json
import argparse
from pyecharts.charts.chart import ThreeAxisChart
import pyecharts.options as opts
from pyecharts.charts import Line, Grid, Tab


class ViewMem():
    _save_dir = None
    _hzmm_log_path = None
    _caller_id_path = None
    _caller_name_path = None
    _caller_name_json = None
    _thread_id_path = None
    _thread_name_path = None
    _thread_name_json = None
    _cumulate_json = None
    _hzmm_html_path = None

    cumulate = {}
    _interval = 1
    _last_sec = 0
    caller_name_map = {}
    thread_name_map = {}
    LINES_MAX_NUM_PER_FILE = 4000


    def __init__(self, save_dir, hzmm_log_name) -> None:
        self._save_dir = save_dir
        self._hzmm_log_path = os.path.join(save_dir, "hzmm.log")
        self._caller_id_path = os.path.join(save_dir, "caller.log")
        self._caller_name_path = os.path.join(save_dir, "caller_name.log")
        self._caller_name_json = os.path.join(save_dir, "caller_name.json")
        self._thread_id_path = os.path.join(save_dir, "thread.log")
        self._thread_name_path = os.path.join(save_dir, "thread_name.log")
        self._thread_name_json = os.path.join(save_dir, "thread_name.json")
        self._cumulate_json = os.path.join(save_dir, "cumulate.json")
        self._hzmm_html_path = os.path.join(save_dir, "hzmm.html")

    def read_cumulate(self, model):
        if "new" == model:
            self._read_cumulate_from_log()
        elif "update" == model:
            self._read_cumulate_from_json()
            self._read_cumulate_from_log(self.cumulate['info']['line_num'])
        else:
            print("model error!")
            exit()

    def _read_cumulate_from_log(self, line_num_satrt=0):
        '''
        从 hzmm_log 中读取内存信息，解析并记录，同时按时间间隔（至少1s）来记录，用于生成表格
        '''
        if 0 == line_num_satrt:
            self.cumulate['info'] = {}
            self.cumulate['info']['line_num'] = 0
            self.cumulate['data'] = {}
        cumulate = self.cumulate['data']
        file_index = 0
        line_num = 0
        while True:
            hzmm_log = "hzmm%03d.log" % (file_index)
            hzmm_log_path = os.path.join(self._save_dir, hzmm_log)
            try:
                file = open(hzmm_log_path, 'r')
            except:
                if file_index == 0:
                    return
                else:
                    break
            for line in file:
                if line_num % 10000 == 0:
                    print("> %d" % (line_num))
                line = line.split("|")
                if len(line) < 7:
                    break
                line_num += 1
                if line_num <= line_num_satrt:
                    continue
                if line_num % 500000 == 0:
                    self._write_caller_id_to_disk()
                    self._write_thread_id_to_disk()
                thread_id = line[0] + "_" + line[1]
                sec       = line[2]
                op        = line[3]
                caller    = line[4]
                ptr       = line[5]
                size      = int(line[6])
                if thread_id not in cumulate.keys():
                    cumulate[thread_id] = {}
                if caller not in cumulate[thread_id].keys() and "fre" != op:
                    cumulate[thread_id][caller] = {}
                    cumulate[thread_id][caller]["times"] = 0
                    cumulate[thread_id][caller]["block"] = 0
                    cumulate[thread_id][caller]["bytes"] = 0
                    cumulate[thread_id][caller]["record"] = {}
                    cumulate[thread_id][caller]["sec_list"] = []
                    cumulate[thread_id][caller]["time_list"] = []
                    cumulate[thread_id][caller]["block_list"] = []
                    cumulate[thread_id][caller]["byte_list"] = []
                if "mll" == op:
                    self._malloc(cumulate[thread_id][caller],ptr, size)
                elif "cll" == op:
                    self._calloc(cumulate[thread_id][caller],ptr, size)
                    pass
                elif "rll" == op:
                    self._realloc(cumulate[thread_id][caller], ptr, size)
                    pass
                elif "mml" == op:
                    self._memalign(cumulate[thread_id][caller],ptr, size)
                    pass
                elif "fre" == op:
                    self._free(ptr)
                self._record_cumulate(sec)
            file.close()
            file_index += 1
            if line_num % self.LINES_MAX_NUM_PER_FILE != 0:
                break
        self.cumulate['info']['line_num'] = line_num
        self._write_cumulate_to_json()

    def _malloc(self, CALLER, ptr, size):
        CALLER["times"] += 1
        CALLER["block"] += 1
        CALLER["bytes"] += size
        if ptr not in CALLER["record"].keys():
            CALLER["record"][ptr] = 0
        CALLER["record"][ptr] += size

    def _calloc(self, CALLER, ptr, size):
        CALLER["times"] += 1
        CALLER["block"] += 1
        CALLER["bytes"] += size
        if ptr not in CALLER["record"].keys():
            CALLER["record"][ptr] = 0
        CALLER["record"][ptr] += size

    def _realloc(self, CALLER, ptr, size):
        CALLER["times"] += 1
        CALLER["block"] += 1
        CALLER["bytes"] += size
        if ptr not in CALLER["record"].keys():
            CALLER["record"][ptr] = 0
        CALLER["record"][ptr] += size

    def _memalign(self, CALLER, ptr, size):
        CALLER["times"] += 1
        CALLER["block"] += 1
        CALLER["bytes"] += size
        if ptr not in CALLER["record"].keys():
            CALLER["record"][ptr] = 0
        CALLER["record"][ptr] += size

    def _free(self, ptr):
        cumulate = self.cumulate['data']
        for i in cumulate.keys():
            thread_id = cumulate[i]
            for j in thread_id.keys():
                caller = thread_id[j]
                need_del_key = None
                for k in caller["record"].keys():
                    if k == ptr:
                        caller["block"] -= 1
                        caller["bytes"] -= caller["record"][k]
                        #need_del_key.append(k)
                        need_del_key = k
                if need_del_key != None:
                    caller["record"].pop(need_del_key)
                    return
                #for k in need_del_key:
                #    caller["record"].pop(k)

    def _record_cumulate(self, sec):
        if self._last_sec == 0:
            self._last_sec = sec
        cumulate = self.cumulate['data']
        if int(sec) - int(self._last_sec) > self._interval:
            self._last_sec = sec
            for i in cumulate.keys():
                thread_id = cumulate[i]
                for j in thread_id.keys():
                    caller = thread_id[j]
                    caller["sec_list"].append(sec)
                    caller["time_list"].append(caller["times"])
                    caller["block_list"].append(caller["block"])
                    caller["byte_list"].append(caller["bytes"])

    def _get_caller_and_thread_name(self):
        if len(self.caller_name_map.keys()) != 0 and \
            len(self.thread_name_map.keys()) != 0:
            return

        self._write_caller_id_to_disk()
        self._write_thread_id_to_disk()
        print("Get caller name and thread name. Wait 3 second..")
        time.sleep(1)
        print("Get caller name and thread name. Wait 2 second..")
        time.sleep(1)
        print("Get caller name and thread name. Wait 1 second..")
        time.sleep(1)

        if 0 == self._request_caller_name():
            self._write_caller_name_to_json()
            os.remove(self._caller_name_path)
        else:
            self._read_caller_name_from_json()

        if 0 == self._request_thread_name():
            self._write_thread_name_to_json()
            #os.remove(self._thread_name_path)
        else:
            self._read_thread_name_from_json()

    def _write_caller_name_to_json(self):
        try:
            f = open(self._caller_name_json, "w", encoding='utf-8')
        except:
            return
        json.dump(self.caller_name_map, f, ensure_ascii=False, sort_keys=False, \
            indent=2, separators=(',', ':'))

    def _read_caller_name_from_json(self):
        try:
            f = open(self._caller_name_json, 'r')
        except:
            return
        self.caller_name_map = json.load(f)

    def _write_thread_name_to_json(self):
        try:
            f = open(self._thread_name_json, "w", encoding='utf-8')
        except:
            return
        json.dump(self.thread_name_map, f, ensure_ascii=False, sort_keys=False, \
            indent=2, separators=(',', ':'))

    def _read_thread_name_from_json(self):
        try:
            f = open(self._thread_name_json, 'r')
        except:
            return
        self.thread_name_map = json.load(f)

    def _write_cumulate_to_json(self):
        try:
            f = open(self._cumulate_json, "w", encoding='utf-8')
        except:
            return
        json.dump(self.cumulate, f, ensure_ascii=False, sort_keys=False, \
            indent=2, separators=(',', ':'))

    def _read_cumulate_from_json(self):
        try:
            f = open(self._cumulate_json, 'r')
        except:
            return
        self.cumulate = json.load(f)


    def _write_caller_id_to_disk(self):
        f = open(self._caller_id_path, "w")
        cumulate = self.cumulate['data']
        for i in cumulate.keys():
            thread_id = cumulate[i]
            for j in thread_id.keys():
                f.write(j + "\n")
        f.write("#")
        f.close()

    def _request_caller_name(self):
        caller_id_list = []
        caller_name_list = []
        try:
            file = open(self._caller_name_path, 'r')
        except:
            return -1
        for line in file:
            line = line.split("\n")[0]
            if "|" in line:
                line = line.split("|")[1]
                caller_id_list.append(line)
            else:
                caller_name_list.append(line)
        if len(caller_id_list) != len(caller_name_list):
            return -1
        for i in range(len(caller_id_list)):
            self.caller_name_map[caller_id_list[i]] = caller_name_list[i]
        return 0

    def _write_thread_id_to_disk(self):
        f = open(self._thread_id_path, "w")
        cumulate = self.cumulate['data']
        for i in cumulate.keys():
            id = i.split("_")
            pid = id[0]
            tid = id[1]
            f.write(pid + " " + tid + "\n")
        f.write("#")
        f.close()

    def _request_thread_name(self):
        thread_id_list = []
        thread_name_list = []
        try:
            file = open(self._thread_name_path, 'r')
        except:
            return -1
        for line in file:
            line = line.split("\n")[0]
            line = line.split("|")
            if len(line) < 2:
                break
            thread_id_list.append(line[0])
            thread_name_list.append(line[1])
        if len(thread_id_list) != len(thread_name_list):
            return -1
        for i in range(len(thread_id_list)):
            self.thread_name_map[thread_id_list[i]] = thread_name_list[i]
        return 0

    def _get_line(self, caller_name, CALLER):
        line = Line(
            init_opts=opts.InitOpts(
                #js_host="./",
            )
        )
        for key in CALLER.keys():
            if key == "sec_list":
                line.add_xaxis(CALLER[key])
            elif key == "time_list":
                line.add_yaxis(key, CALLER[key])
            elif key == "block_list":
                line.add_yaxis(key, CALLER[key])
            elif key == "byte_list":
                line.add_yaxis(key, CALLER[key])
        line.set_global_opts(
            yaxis_opts=opts.AxisOpts(is_scale=True, name=caller_name),
            xaxis_opts=opts.AxisOpts(is_scale=True),
            datazoom_opts=[opts.DataZoomOpts(type_="inside",
                range_start=0,
                range_end=100,
                )],
        )
        return line

    def _get_grid(self, THREAD):
        caller_num = len(THREAD.keys())
        height = caller_num * 500
        grid = Grid(
            init_opts=opts.InitOpts(
                #js_host='./',
                width='1900px',
                height='%dpx'%(height),
                animation_opts=opts.AnimationOpts(animation=False),
            )
        )
        for i, key in enumerate(THREAD.keys()):
            caller = THREAD[key]
            if key in self.caller_name_map.keys():
                line = self._get_line(self.caller_name_map[key], THREAD[key])
            else:
                line = self._get_line(key, THREAD[key])
            grid.add(line, grid_opts=opts.GridOpts(pos_left='120px',
                                                   pos_top='%dpx'%(500*i+30),
                                                   height='400px'))
        return grid

    def creat_html(self):
        self._get_caller_and_thread_name()
        tab = Tab()#(js_host='./')
        cumulate = self.cumulate['data']
        for i in cumulate.keys():
            thread_id = cumulate[i]
            grid = self._get_grid(thread_id)
            if i in self.thread_name_map.keys():
                tab.add(grid, self.thread_name_map[i])
            else:
                tab.add(grid, i)
        tab.render(self._hzmm_html_path)

    def print_cumulate(self):
        self._get_caller_and_thread_name()
        print("\nthread\tcaller\t\ttimes\tblocks\tbytes\tthread name")
        cumulate = self.cumulate['data']
        for i in cumulate.keys():
            if i in self.thread_name_map.keys():
                print(self.thread_name_map[i])
            else:
                print(i)
            thread_id = cumulate[i]
            for j in thread_id.keys():
                caller = thread_id[j]
                if j in self.caller_name_map.keys():
                    print("\t%s:\t%d\t%d\t%d\t%s" % (j, caller["times"], caller["block"], \
                        caller["bytes"], self.caller_name_map[j]))
                else:
                    print("\t%s:\t%d\t%d\t%d\t" % (j, caller["times"], caller["block"], \
                        caller["bytes"]))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--model', default="None")
    args = parser.parse_args()
    if args.model == "None":
        print("Usage:")
        print("  python view_mem.py -m update   # '-m new' ")
        exit()
    vm = ViewMem(save_dir="mmd", hzmm_log_name="hzmm.log")
    vm.LINES_MAX_NUM_PER_FILE = 16000000 # 和 hzmm.c 里的保持一致
    vm.read_cumulate(model=args.model)
    vm.print_cumulate()
    vm.creat_html()
